import os
import tensorflow as tf
from data_statistic import statistic


BATCH_SIZE = 64


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])


def load_image(path, size=(32, 32)):
    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img)  # 注意此处为jpeg格式
    img = tf.image.resize(img, size) / 255.0
    return img


def useModel(ruler, path):
    # 对path下所有文件使用模型
    model = tf.keras.models.load_model('model/tf_model_savedmodel')
    path = "predict/"
    positive_results = []
    negative_results = []
    for root, dirs, files in os.walk(path):
        for file in files:
            path = "predict/" + file
            data = tf.data.Dataset.list_files(path) \
                .map(load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE) \
                .batch(BATCH_SIZE) \
                .prefetch(tf.data.experimental.AUTOTUNE)

            # 加载模型
            res = [model.predict_on_batch(x) for x in data.take(1)][0]
            for i in res:
                if i > 0.5:
                    area, length, width, angle = statistic(path, ruler)
                    result = [file, area, length, width, angle]
                    positive_results.append(result)

                else:
                    negative_results.append(file)
    return positive_results,negative_results
